# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: yuxuan

import collections


class CoolDownLimiter(object):
  def __init__(self, cooldown_ns, cooldown_limit=1):
    self._cooldown_ns = cooldown_ns
    self._cooldown_limit = cooldown_limit
    self._trans_q = collections.deque()
    self._trans_d = dict()

  def add(self, timestamp, key):
    self._update(timestamp)
    count = self._trans_d.get(key, 0)
    if count < self._cooldown_limit:
      self._trans_d[key] = count + 1
      self._trans_q.append([timestamp, key])
      return True
    return False

  def _update(self, timestamp):
    while self._trans_q and timestamp - self._trans_q[0][0] > self._cooldown_ns:
      key = self._trans_q[0][1]
      count = self._trans_d[key]
      count -= 1
      if count == 0:
        del self._trans_d[key]
      else:
        self._trans_d[key] = count
      self._trans_q.popleft()

  def in_cooldown(self, timestamp, key):
    self._update(timestamp)
    return self._trans_d.get(key, 0) >= self._cooldown_limit
